import os
import subprocess
import time
import signal
import rosgraph
import rospy
import std_msgs.msg as std_msgs
from std_msgs.msg import Bool
from nav_msgs.msg import Odometry
import argparse as argparse
import sys
from tqdm import tqdm
import asyncio
import websockets
import json
from rospy.msg import AnyMsg
from websockets.sync.client import connect

args = argparse.ArgumentParser()
args.add_argument("--finish-line", type=float, default=40)
args.add_argument("--trigger-topic", type=str, default="/bfctrl/cmd")
args.add_argument("--odom-topic", type=str, default="/bfctrl/local_odom")
args.add_argument(
    "--collision-topic", type=str, default="/airsim_node/drone_1/collision"
)
args.add_argument("--max-iter", type=int, default=20)
args.add_argument("--max-time", type=int, default=60)
args.add_argument("--shfile", type=str, default="tools/avd_mpc.sh")
args.add_argument("--output_dir", type=str, default="./output")
args.add_argument("--test-args", nargs="+", type=str, default=[])
args.add_argument("--tags", nargs="+", type=str, default=["2mps"])
args.add_argument("--save-ply", action="store_true", help="Save ply file")
args = args.parse_args()
output_dir = args.output_dir
output_dir = os.path.abspath(output_dir)


def check_roscore():
    if rosgraph.is_master_online():
        return True
    else:
        return False


class RosMonitior:
    def __init__(self) -> None:
        rospy.init_node("auto_test", anonymous=True)
        self.start_time = time.time()
        self.statue = "Init"
        self.reset_pub = rospy.Publisher("/auto_test/reset", Bool, queue_size=1)
        self.triggerSub = rospy.Subscriber(
            args.trigger_topic, AnyMsg, self.TriggerCallback
        )
        self.video_name = None
        self.video_name_sub = rospy.Subscriber(
            "/video_names", std_msgs.String, self.VideoNameCallback
        )
        self.is_collision = False
        self.collisionSub = rospy.Subscriber(
            args.collision_topic, Bool, self.CollisionCallback
        )
        self.local_odom = None
        self.local_odomSub = rospy.Subscriber(
            args.odom_topic,
            Odometry,
            self.LocalOdomCallback,
        )
        self.lap_count = 0
        self.lap_count_sub = rospy.Subscriber(
            "/lap_count", std_msgs.Int32, self.LapCountCallback
        )

    def TriggerCallback(self, msg: AnyMsg):
        self.statue = "CMD"

    def CollisionCallback(self, msg: std_msgs.Bool):
        if msg.data == True and self.statue == "CMD":
            self.is_collision = True

    def VideoNameCallback(self, msg: std_msgs.String):
        self.video_name = msg.data

    def LocalOdomCallback(self, msg: Odometry):
        self.local_odom = msg

    def LapCountCallback(self, msg: std_msgs.Int32):
        self.lap_count = msg.data

    def IsFinished(self):
        duration = time.time() - self.start_time

        if self.local_odom is not None and (
            self.local_odom.pose.pose.position.x > args.finish_line
            or self.lap_count > 0
        ):
            return True, "success"
        elif self.is_collision:
            return True, "collision"
        elif duration > args.max_time:
            return True, "timeout"
        else:
            return False, "running"

    def Reset(self):
        self.reset_pub.publish(Bool(True))


def signal_handler(sig, frame):
    print("Interrupted! Cleaning up...")
    try:
        report_file.write(f"## Overview\n\n")
        report_file.write(f"**Total Iteration**: {total_count}\n\n")
        report_file.write(f"**Scuccess Iteration**: {success_count}\n\n")
        report_file.write(f"**Success rate**: {success_count}/{total_count}\n\n")
        report_file.flush()
    except:
        print("Variable 'report_file' does not exist.")
    try:
        os.killpg(os.getpgid(cmd_subprocess.pid), signal.SIGTERM)
        cmd_subprocess.wait()
    except:
        print("Variable 'cmd_subprocess' does not exist.")
    try:
        os.killpg(os.getpgid(roscore_subprocess.pid), signal.SIGTERM)
        roscore_subprocess.wait()
    except:
        print("Variable 'roscore_subprocess' does not exist.")
        sys.exit(0)


def send_video_info(video_path, video_tags, bags_path=None):
    async def send():
        uri = "ws://localhost:8765"
        async with websockets.connect(uri, ping_timeout=0.5) as websocket:
            message = {"type":"update_video_infos","video_path": video_path, "video_tags": video_tags, "bag_path": bags_path}
            if bags_path is not None:
                message["bags_path"] = bags_path
            await websocket.send(json.dumps(message))
            response = await websocket.recv()

    try:
        asyncio.get_event_loop().run_until_complete(send())
    except:
        print("Failed to send video info")

def save_ply(save_path):
    uri = "ws://localhost:8765"
    try:
        # Establish a blocking connection
        with connect(uri, open_timeout=60.0, close_timeout=60.0) as websocket:
            # Prepare and send the message
            message = {"type": "get_ply", "save_path": save_path}
            websocket.send(json.dumps(message))            
            # Receive and process the response
            response = websocket.recv()

    except Exception as e:
        print("Failed to save ply:", e)

def reset_scene():
    uri = "ws://localhost:8765"
    try:
        # Establish a blocking connection
        with connect(uri, open_timeout=60.0, close_timeout=60.0) as websocket:
            # Prepare and send the message
            message = {"type": "reset"}
            websocket.send(json.dumps(message))            
            # Receive and process the response
            response = websocket.recv()

    except Exception as e:
        print("Failed to reset the scene:", e)

def reset_random_seed():
    uri = "ws://localhost:8765"
    try:
        # Establish a blocking connection
        with connect(uri, open_timeout=60.0, close_timeout=60.0) as websocket:
            # Prepare and send the message
            message = {"type": "reset_random_seed"}
            websocket.send(json.dumps(message))            
            # Receive and process the response
            response = websocket.recv()

    except Exception as e:
        print("Failed to reset random seed:", e)

def generate_test_id():
    time_id = time.strftime("%Y%m%d%H%M", time.localtime())
    time_id = "AT"+time_id[2:]
    return time_id

def find_bag_file_in_folder(folder):
    bag_files = None
    for root, dirs, files in os.walk(folder):
        for file in files:
            if file.endswith(".bag") or file.endswith(".bag.active"):
                if file.endswith(".active"):
                    file = file[:-7]
                bag_files = os.path.join(root, file)
                break
    return bag_files

signal.signal(signal.SIGINT, signal_handler)
curdate = time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())
sheel_script = args.shfile
sheel_base = os.path.basename(sheel_script).split(".")[0]
out_base = f"{sheel_base}"
for tag in args.tags:
    out_base += f"-{tag}"
out_base += f"-{curdate}"
roscore_is_running = check_roscore()
if not roscore_is_running:
    print("roscore is not running, start roscore")
    command = "roscore"
    roscore_subprocess = subprocess.Popen(command, shell=True, preexec_fn=os.setsid)
    time.sleep(5)
max_iter = args.max_iter
test_id = generate_test_id()
success_count = 0
total_count = 0
pbar = tqdm(range(max_iter))
report_file = os.path.join(args.output_dir, out_base, "report.md")
if not os.path.exists(os.path.join(args.output_dir, out_base)):
    os.makedirs(os.path.join(args.output_dir, out_base))
report_file = open(report_file, "w")
report_file.write(f"# Test Report (id: {test_id})\n\n")
report_file.write(f"**Test Script**: {sheel_script}\n\n")
report_file.write(f"**Test Args**: {' '.join(args.test_args)}\n\n")
report_file.write(f"**Tags**: {', '.join(args.tags)}\n\n")
report_file.write("## Results:\n")
report_file.write("| Iteration | Time | Duration | Result |\n")
report_file.write("| --- | --- | --- | --- |\n")
report_file.flush()
reset_random_seed()
for iter in pbar:
    reset_scene()
    bag_record_dir = os.path.join(args.output_dir, out_base, "bags", f"iter_{iter}")
    bag_record_dir = os.path.abspath(bag_record_dir)
    if not os.path.exists(bag_record_dir):
        os.makedirs(bag_record_dir)
    command = "bash " + sheel_script + " " + str(bag_record_dir) + " "
    for test_arg in args.test_args:
        command += f" {test_arg}" + " "
    log_file = os.path.join(bag_record_dir, "log.txt")
    log_file = open(log_file, "w")
    cmd_subprocess = subprocess.Popen(
        command,
        shell=True,
        preexec_fn=os.setsid,
        stdout=log_file,
        stderr=log_file,
    )
    state_monitor = RosMonitior()
    is_finished, result = state_monitor.IsFinished()
    start_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    start_timestamp = time.time()
    while not is_finished:
        is_finished, result = state_monitor.IsFinished()
        time.sleep(0.1)
    end_timestamp = time.time()
    duration = end_timestamp - start_timestamp
    os.killpg(os.getpgid(cmd_subprocess.pid), signal.SIGINT)
    if state_monitor.video_name is not None:
        tags = []
        for tag in args.tags:
            tags.append(tag)
        tags.append(f"{sheel_base}")
        tags.append(result)
        tags.append(f"iter{iter}")
        tags.append(f"{test_id}")
        bag_file = find_bag_file_in_folder(bag_record_dir)
        send_video_info(state_monitor.video_name, tags, bag_file)
        state_monitor.video_name = None
    if result == "success":
        success_count += 1
    total_count += 1
    pbar.set_description(f"sr: {success_count}/{iter + 1}")
    report_file.write(f"| {iter + 1} | {start_time} | {duration} | {result} |\n")
    report_file.flush()
    cmd_subprocess.wait()
    if args.save_ply:
        save_ply(bag_record_dir)
    time.sleep(1)

report_file.write(f"## Overview\n\n")
report_file.write(f"**Total Iteration**: {max_iter}\n\n")
report_file.write(f"**Scuccess Iteration**: {success_count}\n\n")
report_file.write(f"**Success rate**: {success_count}/{max_iter}\n\n")
report_file.flush()
try:
    os.killpg(os.getpgid(cmd_subprocess.pid), signal.SIGINT)
    cmd_subprocess.wait()
except:
    print("Variable 'cmd_subprocess' does not exist.")
try:
    os.killpg(os.getpgid(roscore_subprocess.pid), signal.SIGINT)
    roscore_subprocess.wait()
except:
    print("Variable 'roscore_subprocess' does not exist.")